[PERF] MXFP4 (a4w4) MoE backend for gfx950 by ColorsWind · Pull Request #3470 · ROCm/aiter

ColorsWind · 2026-06-01T15:52:40Z

[PERF] mxfp4_moe — a native MXFP4 (a4w4) MoE backend for gfx950

Motivation

Kimi-K2.5 / DeepSeek-R1 serve MoE FFNs with MXFP4 (a4w4) weights. The existing FlyDSL (MLIR) a4w4 MoE path leaves performance on the table across the whole token range — it is launch/latency-bound at decode-size M and not fully HBM-bound at prefill-size M. This PR adds mxfp4_moe, a native pure-HIP MXFP4 MoE backend that is faster than FlyDSL at every M bucket on MI350X.

Technical Details

The backend is two JIT modules plus a fused_moe integration:

module_moe_mxfp4_aux — token sort, mxfp4 quant (bf16 → fp4 + e8m0, group=32), scale sort/shuffle, and the weighted scatter-reduce.
module_moe_mxfp4_gemm — the two a4w4 GEMMs (gemm1 fused-SwiGLU, gemm2 down-proj) with codegen'd kernel-name dispatch (gen_instances.py) and MFMA f4f4 microkernels.
fused_moe wiring — mxfp4_moe_* kernel names route to an _mxfp4_moe_run pipeline via a MOEMetadata hook; per-M kernel selection comes from a shuffle_kind-tagged tuned CSV, so the backend ships its own tuned rows next to the default tuning without dedup conflicts.

Several pipelines, dispatched per token-count M

The optimal fusion/accumulation strategy flips between decode-size and prefill-size M, so we designed several pipelines and pick one per M from the tuned CSV:

inline-quant + atomic — small M (BM=16): sort → quant-fused-gemm1 → gemm2 (atomic-add reduce). Quant is folded into gemm1 and gemm2 atomic-adds straight into the output, removing the per-sorted-row partial buffer and the separate quant / scatter passes. (The gemm2 B-load additionally has a per-M-tuned cached vs non-temporal choice.) Wins when M ≲ BM and the path is launch/latency-bound.
three-stage — mid-to-large M (BM=32 / 128): sort → quant → sort-scale → gemm1 → gemm2 [→ scatter-reduce]. Prologue runs as separate full-throughput kernels. At mid M (BM=32) gemm2 atomic-accumulates straight into the output (no separate reduce); at large M (BM=128) gemm2 writes per-sorted-row flat_out and a dedicated non-atomic scatter-reduce does the weighted sum — avoiding the atomic contention that grows with M. HBM-bound regime.

The core tradeoff: fusion + atomics win when operator shapes are complementary and launch overhead dominates (small M); at large M they break down — the same token is read by multiple experts (duplicated quant work) and atomic contention grows — so the separate, non-atomic path wins.

On top of the large-M non-atomic path, an optional mxfp4-intermediate variant (CSV-gated, enabled on the largest prefill buckets) stages gemm2's flat_out as packed fp4 + e8m0 (mxfp4-out epilog) and reads it back with scatter_reduce_q, cutting the reduce's HBM traffic ~3.8×.

Test Plan

Machine / shape

AMD Instinct MI350X (gfx950)
Kimi-K2.5, TP=4: NE=385, H=7168, D_INTER=512, TOPK=9
Routing: 1 shared expert (always selected) + 8 routed experts, softmax with a mild per-expert popularity skew (skew=0.5).
Both backends fed the same per-1x32 MXFP4 weights (quantized once).
FlyDSL baseline: ROCm/aiter @ ca781d81.

Cache / timing

Correctness: end-to-end pipeline vs a full-precision bf16 ground truth.
Performance: end-to-end MoE wall-clock (perf_counter + synchronize), median of 30 iters, warm cache (weights resident in L2; no per-iter flush). mxfp4_moe is driven through the production fused_moe entry — per-M kernel dispatch comes straight from the tuned CSV (including the mxfp4-intermediate path on the largest buckets), so the numbers carry real host-side dispatch overhead; FlyDSL is timed at its kernel stages (best block_m reported). Same routing and fp4 weights for both — if anything this is conservative for mxfp4_moe at small M.

Correctness

End-to-end cos(kernel, bf16_truth) is 0.975–0.978 across the full M range and never worse than the fp4 reference — all pass.

Performance (end-to-end, mxfp4_moe vs FlyDSL)

M	mxfp4_moe (µs)	flydsl (µs)	speedup	cos(mxfp4_moe, flydsl)
4	99.3	178.3	1.80×	0.985
8	116.2	201.7	1.74×	0.987
16	153.8	233.6	1.52×	0.987
32	204.5	300.3	1.47×	0.987
64	254.8	342.8	1.35×	0.987
128	304.2	412.4	1.36×	0.987
256	363.3	471.7	1.30×	0.987
4096	999.7	1115.2	1.12×	0.987
8192	1399.0	1785.9	1.28×	0.981
16384	2250.5	3176.7	1.41×	0.981

mxfp4_moe is faster than FlyDSL across the entire token range (1.12×–1.80×). At the largest buckets (M ≥ 8192) the CSV switches to the mxfp4-intermediate path; its cos vs FlyDSL drops to ~0.981 (the inherent loss of the before-sum fp4 quant) while still pulling ahead in latency.

Add module_moe_mxfp4_aux for the MXFP4 (a4w4) MoE backend: fused token sort+quant, 3-stage sort, scale sort/shuffle, and weighted bf16 scatter-reduce (gfx950, Kimi-K2.5 / DSR shapes).

Add module_moe_mxfp4_gemm: the two MXFP4 (a4w4) MoE GEMMs (gemm1 fused SwiGLU, gemm2 down-proj) with codegen'd kernel-name dispatch (gen_instances.py), MFMA f4f4 microkernels, epilogs, non-temporal B-load dispatch, and XCD-swizzle interface (gfx950, Kimi-K2.5 / DSR).

Route mxfp4_moe_* kernel names to the _mxfp4_moe_run pipeline (sort/quant -> gemm1 -> gemm2 -> scatter_reduce) via a MOEMetadata pipeline hook, with shuffle_kind-tagged CSV lookup so the backend ships its own tuned rows alongside the default set. Add Kimi-K2.5 tuned CSV.

Add the optional mxfp4-intermediate path: gemm2 stages flat_out as packed fp4 + e8m0 (mxfp4-out epilog) and scatter_reduce_q reads it back as mxfp4, cutting the reduce's HBM traffic ~3.8x. Gated by a _MXFP4OUT tag in the tuned CSV (large-M buckets only, where the reduce win beats the gemm2 epilog overhead). bf16-intermediate stays default.

github-actions · 2026-06-01T15:53:03Z

🏷️ CI Guide

Runs automatically on every PR:

✅ Pre-checks (submodule verification, code formatting)
✅ Aiter op tests (gfx942 + gfx950)
✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
`ci:triton-300x`	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
`ci:sglang`	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
`ci:atom`	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
`ci:atom_full`	ATOM accuracy suite for PR and main models from ATOM `models_accuracy.json`
`ci:vllm`	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
`ci:all`	All standard extended tests (excludes `ci:atom_full`)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 3470 --add-label <label>

benenzhu · 2026-06-02T02:50:19Z

Performance bench script: https://gist.github.com/ftyghome/088870c544693da57258d22968573620

coderfeli · 2026-06-03T03:58:03Z

@benenzhu it is launch/latency-bound' in your bench script you use s = time.perf_counter() and then time.perf_counter() - s. It counts all cpu overhead. The real run in e2e uses cudagraph that will not count.

ftyghome · 2026-06-03T08:05:19Z

@benenzhu it is launch/latency-bound' in your bench script you use s = time.perf_counter() and then time.perf_counter() - s. It counts all cpu overhead. The real run in e2e uses cudagraph that will not count.

Hi Felix,

Thanks for the review!
We just updated the microbench script to capture 60 runs in a cudagraph. The new perf data is listed below:

M	mxfp4 us	flydsl us	flydsl/mxfp4	cos
4	43.5	69.9	1.61x	0.9845
8	62.2	96.4	1.55x	0.9868
16	101.6	132.9	1.31x	0.9874
32	151.3	190.2	1.26x	0.9868
64	206.8	236.6	1.14x	0.9874
128	256.3	305.8	1.19x	0.9875
256	310.2	341.0	1.10x	0.9871
4096	977.4	1025.6	1.05x	0.9875
8192	1381.3	1718.0	1.24x	0.9812
16384	2254.4	3133.8	1.39x	0.9811

lalala-sh · 2026-06-03T09:38:20Z

@benenzhu it is launch/latency-bound' in your bench script you use s = time.perf_counter() and then time.perf_counter() - s. It counts all cpu overhead. The real run in e2e uses cudagraph that will not count.

Hi Felix,

Thanks for the review! We just updated the microbench script to capture 60 runs in a cudagraph. The new perf data is listed below:

M mxfp4 us flydsl us flydsl/mxfp4 cos
4 43.5 69.9 1.61x 0.9845
8 62.2 96.4 1.55x 0.9868
16 101.6 132.9 1.31x 0.9874
32 151.3 190.2 1.26x 0.9868
64 206.8 236.6 1.14x 0.9874
128 256.3 305.8 1.19x 0.9875
256 310.2 341.0 1.10x 0.9871
4096 977.4 1025.6 1.05x 0.9875
8192 1381.3 1718.0 1.24x 0.9812
16384 2254.4 3133.8 1.39x 0.9811

Hello @ftyghome, Could you please update the perf compare again? Collect the flydsl date with fused_moe interface, so that flydsl will be called through the moe tuner too.

benenzhu · 2026-06-03T11:07:21Z

@benenzhu it is launch/latency-bound' in your bench script you use s = time.perf_counter() and then time.perf_counter() - s. It counts all cpu overhead. The real run in e2e uses cudagraph that will not count.

Hi Felix,
Thanks for the review! We just updated the microbench script to capture 60 runs in a cudagraph. The new perf data is listed below:
M mxfp4 us flydsl us flydsl/mxfp4 cos
4 43.5 69.9 1.61x 0.9845
8 62.2 96.4 1.55x 0.9868
16 101.6 132.9 1.31x 0.9874
32 151.3 190.2 1.26x 0.9868
64 206.8 236.6 1.14x 0.9874
128 256.3 305.8 1.19x 0.9875
256 310.2 341.0 1.10x 0.9871
4096 977.4 1025.6 1.05x 0.9875
8192 1381.3 1718.0 1.24x 0.9812
16384 2254.4 3133.8 1.39x 0.9811

Hello @ftyghome, Could you please update the perf compare again? Collect the flydsl date with fused_moe interface, so that flydsl will be called through the moe tuner too.

Hi updated: @lalala-sh
Script: https://gist.github.com/benenzhu/ea38af8e5606ed5596abc1890e56a5a1
Change by ai, not verify the correctness.

M	mxfp4 us	bm	flydsl us	flydsl/mxfp4	mx·fly	mx·ref	fly·ref
4	41.9	16	59.2	1.41x	0.9845	0.9533	0.9482
8	61.2	16	74.6	1.22x	0.9868	0.9524	0.9485
16	100.3	16	122.1	1.22x	0.9876	0.9538	0.9514
32	147.8	16	165.0	1.12x	0.9868	0.9548	0.9520
64	203.5	16	213.9	1.05x	0.9874	0.9555	0.9529
128	254.2	16	263.8	1.04x	0.9875	0.9562	0.9536
256	304.1	32	306.4	1.01x	0.9871	0.9561	0.9538
8192	1137.3	128	1359.0	1.19x	0.9811	0.9505	0.9541
16384	1821.7	128	2325.2	1.28x	0.9810	0.9506	0.9541

coderfeli · 2026-06-05T12:56:49Z

+        }
+        if (NE == 257) {
+            if (MB == 32) { LAUNCH(257, 9, 32, 7168); return; }
+        }


why so many constraints here? Could you make it more general ? I think the kernels don't have these limitation?

It's the deepseek & kimi shapes. only tested on this.

coderfeli · 2026-06-05T12:57:02Z

+            }
+            if (NE == 257) {
+                if (MB == 16) { LAUNCH_IQ_ZI(257, 9, 16, 7168); return; }
+            }


coderfeli · 2026-06-05T12:57:09Z

+            if (NE == 257) {
+                if (MB == 16) { LAUNCH_IQ(257, 9, 16, 7168); return; }
+            }
+        }


coderfeli · 2026-06-05T12:58:38Z

+        }
+        if (NE == 257) {
+            if (MB == 32)  { LAUNCH(257, 9, 32,  7168); return; }
+            if (MB == 128) { LAUNCH(257, 9, 128, 7168); return; }


Also make quant a general usable kernel?

The general case is already covered by dynamic_mxfp4_quant in aiter, which work on arbitrary shapes. This mxfp4_moe_quant is a fused, shape-specialized variant (compile-time D_HIDDEN for vectorization, plus the optional fused zero-init) that only makes sense inside this MoE pipeline.

If you need it for more model shapes, the cheap path is to add template dispatch cases (keeps the compile-time perf) rather than making it fully runtime-general. Happy to do that as a follow-up if there's a concrete shape you need.

coderfeli · 2026-06-05T13:04:05Z

+__launch_bounds__(256,
+                  is_nonatomic_v<kEpilog> ? 1 :
+                  ((BM == 16) ? 4 : 2))
+kernel(


ROCm/FlyDSL#660 for gemm2, we try to impl a flydsl pipeline that's exactly the same as this one. Would you mind transfer the same style and infra for better long term maintaining?

Sure. We also have a FlyDSL version of gemm1, and most of the performance already matches our HIP implementation in our tests (except M=4 case). We will upload it soon.

It would be great to sync with the ongoing gemm1/gemm2 FlyDSL work. We would be happy to work together on this and align the performance with our HIP implementation while moving toward FlyDSL.

Yes. Thanks a lot. HIP backend is also great. But currently there are ck in hip, flydsl, which we are trying to unify. Current PR add another new one, we hope to maintain in one infra for all optimization if possible.

coderfeli · 2026-06-05T13:05:02Z

+
+    if (D_HIDDEN == 7168 && D_INTER == 512) {
+        if (NE == 385) { LAUNCH_G2_MXFP4(385, 512, 7168); return; }
+        if (NE == 257) { LAUNCH_G2_MXFP4(257, 512, 7168); return; }


hard code shapes. also couldn't be extended.

coderfeli · 2026-06-05T13:05:29Z

+        sorted_weights.data_ptr<float>(),                                                        \
+        reinterpret_cast<__hip_bfloat16*>(out.data_ptr()))
+
+    if (D_HIDDEN == 7168 && TOPK == 9) {


coderfeli · 2026-06-05T13:06:19Z

+    const bool nt_hints = (MB >= 128);
+
+#define LAUNCH_Q(D_HIDDEN_, TOPK_, NT_)                                                           \
+    aiter::mxfp4_moe::moe_scatter_reduce::launch_mxfp4<                                           \


All input is quanted to fp4. It's ok for gsm8k but not in our many other real tests. could you also add flags to enable ?

Made it opt-in: set AITER_MXFP4_INTERMEDIATE=1 to enable the fp4 path, otherwise it falls back to the accurate bf16 reduce by default. So your other evals are safe unless you turn it on. a3c9875

coderfeli · 2026-06-05T13:09:02Z

+    const buffer_rsrc_t B_q_rsrc =
+        make_buffer_rsrc(B_q,
+            (uint32_t)((long long)NUM_EXPERTS * N_OUT * K_HALF * sizeof(__hip_fp4x2_storage_t)));
+    constexpr int kAS_bound_div = kAtomic ? BM_GRID : 32;


The A_scale buffer-resource bound is 2× too small for BM=64 (atomic path).

A_scale is laid out in fixed 32-row chunks, and the read side addresses it accordingly at line 333:
const int chunk_base = (BM == 16) ? (m_row / BM_GRID) : (m_row / 32);
a_scale_s_base[sub] = (chunk_base + sub) * kAS_per_chunk_dw * 4;
So the chunk-stride divisor is 16 for BM=16 and 32 for everything else — it is not BM. But here the bound divides by kAtomic ? BM_GRID :
32, which only matches the addressing for BM∈{16,32}:
For BM=64 the rsrc extent is (MAX_M/64)*kAS_per_chunk_dw, but addressing reaches (MAX_M/32)*kAS_per_chunk_dw. Once m_row ≥ MAX_M/2, the
hardware buffer-bound check fails the scale loads and returns 0 → e8m0 = 0 → scale 2^-127 → those output rows silently collapse to ~0.
Wrong results, no error.

Good catch. A_scale is addressed in 32-row chunks ((BM==16)?16:32, not BM), but the bound divided by kAtomic ? BM_GRID : 32, so for BM=64 it covers only max_sorted/64 chunks while addressing reaches max_sorted/32. Rows past max_sorted/2 then read out of bounds, the e8m0 comes back 0, and those outputs collapse to ~0.

Fixed by matching the bound to the addressing: (BM==16) ? BM_GRID : 32. It was masked by the old large MAX_M and only surfaced once the bound was tightened to max_sorted. It also went unnoticed because the BM=64 gemm2 kernel is never dispatched (the tuned CSV only uses BM in {16,32,128}, and there is no BM=64 gemm1/sort to feed it).

Fix: bd5ba79, test: 09e0e1d (BM=32 vs BM=64 on identical inputs; fails before, passes after).

coderfeli · 2026-06-05T14:00:46Z

+    (257,  7168,      512,      9),  # DSR
+]
+
+MAX_M = 655360


this is the hard limit of topk * tokens? Then all prefill inputs >72.8K will fail and produce wrong outputs?

You're right that MAX_M is a hard cap on the sorted token count, and inputs past it would produce wrong outputs. It's not really a bug in practice, since MAX_M is set far above any realistic single-call input, so nothing hits it. This change just removes the fixed cap and uses the actual sorted length (M * topk, i.e. sorted_m) for the buffer-resource bounds check, so it scales with the real input instead of relying on a large constant.

See 82e8e0a for the fix and ee60910 for the 128k validation test.

…emm2 kernels Eliminate the MAX_M template parameter from the gemm1 and gemm2 kernels, replacing it with a runtime-calculated max_sorted value. This change enhances flexibility and correctness in buffer resource calculations, ensuring proper handling of varying input sizes. Update related launch functions and code generation scripts accordingly.

…ffer bound Validate the runtime max_sorted fix at 128k tokens via the per-token MoE invariant moe(128k) == concat(16 x moe(8k)): the 8k chunks stay within any bound (trusted reference) while the 128k run exercises the large sorted-row count (>655360) that the old compile-time MAX_M clipped. Uses random weights (self-consistency needs only identical weights + routing across runs) and masks the kernel's inherent FP-reduce nondeterminism by counting a break only where the kernel is self-deterministic yet full != chunked.

The A_scale buffer-resource bound divided by `kAtomic ? BM_GRID : 32`, but the read side addresses A_scale in fixed 32-row chunks (chunk_base = (BM==16) ? m_row/BM_GRID : m_row/32), i.e. the divisor is 16 only for BM=16 and 32 otherwise, never BM. For BM=64 this under-sized the descriptor 2x: addressing reaches max_sorted/32 chunks while the bound only covers max_sorted/64, so every sorted row with m_row >= max_sorted/2 fails the hardware bounds check, the scale load returns 0 (e8m0=0 -> 2^-127), and those output rows silently collapse toward zero. Match the bound divisor to the addressing: (BM==16) ? BM_GRID : 32. This was masked under the old fixed MAX_M (large enough to cover everything) and only surfaced once the bound was tightened to the runtime max_sorted.

Drives mxfp4_moe_gemm2_a4w4 directly and compares the BM=32 atomic kernel (correct bound) against BM=64 (the kernel whose A_scale bound was 2x too small) on identical inputs, which must agree since both address A_scale in 32-row chunks. M_logical/cumsum are chosen so the BM=64 sorted length exceeds the buggy max_sorted/2, placing the upper rows in the clipped region. fp4 partials are forced non-negative so the per-token atomic sum has no cancellation (noise floor ~0), making the clipped-row signal unambiguous. Fails on the buggy bound (7688/8192 rows), passes after the fix (0/8192).

…4_INTERMEDIATE The fp4 gemm2-output path quantizes each expert's output to 4-bit before the topk reduce. That before-sum quant is lossy: fine for gsm8k but it degrades other evals. Make it opt-in via AITER_MXFP4_INTERMEDIATE=1 (off by default), so real evals run the accurate bf16 reduce unless explicitly enabled.

Replace the eight hand-written template-dispatch switches in mxfp4_moe_aux.cu with codegen'd extern "C" instances + string-keyed lookup tables, mirroring the gemm1/gemm2 dispatch. gen_instances.py gains a --target {gemm,aux} flag so a single shared SHAPES list drives both: adding a shape regenerates gemm and aux together. New aux/codegen/mxfp4_moe_aux_dispatch.h holds the per-entry function pointer types plus launch-config constants; module_moe_mxfp4_aux now runs the codegen via blob_gen_cmd (was none). Covers all eight aux entries: sort_quant, sort (threestage / inline_quant / inline_quant+zero_init), quant, sort_scales, scatter_reduce, scatter_reduce_q. Generated instances are byte-identical template instantiations of the prior launches; verified numerically identical to the hand-switch build (old and new show the same FP nondeterminism on the same cancellation rows) and all MB16/MB32/MB128 keys resolve e2e on the Kimi NE=385 shape.

The single full-vs-full noise sample sporadically missed catastrophic- cancellation rows (whose full-vs-chunk difference comes from the differing reduction order, not a bug), causing ~1-in-4 false failures. Estimate the noise floor as a union over several full re-runs (--noise-reps, default 4) and, since full/chunked use different token counts and a same-order probe can never fully capture cross-order cancellation, gate on the COUNT of residual breaks: the max_sorted bug corrupts ~half the rows (~NT/2) while cancellation noise is a handful, so tolerate <0.1% of rows and fail only on a bug-scale fraction. Reliable PASS across repeated runs; still catches the bug with a ~500x margin.

ftyghome · 2026-06-11T07:03:56Z

Hi @coderfeli,

Thanks for all the feedback!

We went through all your inline comments about the hardcoded shapes and constraints, and have made several commits to address them.

Initially, we only tuned and tested the Kimi-K2.5 TP4 shape in detail. The real constraint, however, is just D_HIDDEN % 32 == 0, which comes from the MXFP4 group size. Everything else is a template parameter.

We have now wired the aux kernels into the same CSV codegen flow as gemm1 / gemm2 in 28e1737. With this change, supporting a new shape only requires adding it to the CSV, without further kernel edits.

Could you please take another look when you get a chance? Please let us know if there is anything else we should address.

ftyghome · 2026-06-12T02:44:09Z

Hi @coderfeli,

Based on the latest commits, we added a new end-to-end evaluation for the native HIP MXFP4 a4w4 MoE backend on Kimi-K2.5.

The test compares three configs:

Name	MoE backend	`AITER_MXFP4_INTERMEDIATE`
base	base 2-stage baseline	n/a
MXFP4 intermediate OFF	native HIP MXFP4 a4w4	OFF, bf16 intermediate/reduce
MXFP4 intermediate ON	native HIP MXFP4 a4w4	ON, packed fp4 intermediate from gemm2 to scatter/reduce

Technical detail

The backend uses different pipelines for different token ranges:

small M / decode, BM=16: sort -> quant-fused-gemm1 -> gemm2 (atomic-add reduce)
mid-large M, BM=32 / 128: sort -> quant -> sort-scale -> gemm1 -> gemm2
optional fp4 intermediate path: when AITER_MXFP4_INTERMEDIATE=1, gemm2 stages output packed fp4 intermediate values before scatter/reduce. By default this path is OFF and the backend uses the more accurate bf16 intermediate/reduce path.

Test Plan

Environment:

# 4x gfx950 / MI355X, TP4
# torch / ROCm: 2.10.0+rocm7.2.3

export ATOM_USE_TRITON_MXFP4_BMM=1
export AITER_QUICK_REDUCE_QUANTIZATION=INT4
export ATOM_ENABLE_HIP_MLA_DECODE=1

python -m atom.entrypoints.openai_server --model Kimi-K2.5-MXFP4 \
    -tp 4 --kv_cache_dtype fp8 --max-model-len 10240 \
    --max-num-batched-tokens 16384 --gpu-memory-utilization 0.7 --no-enable_prefix_caching

Accuracy

gsm8k, full 1319 examples, fewshot=5, eval concurrency=64.

Config	strict-match	flexible-extract
base	0.9697	0.9689
MXFP4 intermediate OFF	0.9689	0.9689
MXFP4 intermediate ON	0.9682	0.9674

The fp4 intermediate path only changes gsm8k by about -0.0015 compared with intermediate OFF, which is within the expected evaluation noise.

Performance

Median, Kimi-K2.5 TP4, ISL 8192 / OSL 1024.
Improvement is computed against the base.

Batch	E2E (s)	Impr	Interactivity (tok/s/user)	Impr	TPOT (ms)	Impr	Throughput/GPU (tok/s)	Impr
4	8.74 -> 7.84	+10.3%	106.4 -> 118.6	+11.5%	9.40 -> 8.43	+10.3%	913 -> 1017	+11.4%
8	10.67 -> 9.75	+8.6%	88.5 -> 96.9	+9.5%	11.30 -> 10.32	+8.7%	1516 -> 1659	+9.4%
16	15.09 -> 13.94	+7.6%	61.6 -> 66.8	+8.4%	16.23 -> 14.98	+7.7%	2130 -> 2313	+8.6%
32	21.58 -> 20.14	+6.7%	42.8 -> 45.9	+7.2%	23.34 -> 21.77	+6.7%	2979 -> 3206	+7.6%
64	33.05 -> 31.53	+4.6%	28.4 -> 29.5	+3.9%	35.27 -> 33.92	+3.8%	3860 -> 4033	+4.5%
128	51.30 -> 48.68	+5.1%	19.4 -> 20.4	+5.2%	51.43 -> 48.93	+4.9%	4901 -> 5177	+5.6%
256	92.62 -> 86.41	+6.7%	13.7 -> 14.2	+3.6%	73.18 -> 70.63	+3.5%	4889 -> 5294	+8.3%

Detailed Performance

Interactivity, tok/s/user

Batch	base	MXFP4 intermediate OFF	MXFP4 intermediate ON
4	106.4	118.4	118.6
8	88.5	96.1	96.9
16	61.6	66.3	66.8
32	42.8	45.4	45.9
64	28.4	29.4	29.5
128	19.4	20.1	20.4
256	13.7	13.7	14.2

TPOT, ms

Batch	base	MXFP4 intermediate OFF	MXFP4 intermediate ON
4	9.40	8.45	8.43
8	11.30	10.40	10.32
16	16.23	15.09	14.98
32	23.34	22.01	21.77
64	35.27	34.04	33.92
128	51.43	49.70	48.93
256	73.18	73.12	70.63

Throughput/GPU, tok/s

Batch	base	MXFP4 intermediate OFF	MXFP4 intermediate ON
4	913	1015	1017
8	1516	1647	1659
16	2130	2301	2313
32	2979	3155	3206
64	3860	4005	4033
128	4901	5082	5177
256	4889	5065	5294

E2E latency, s

Batch	base	MXFP4 intermediate OFF	MXFP4 intermediate ON
4	8.74	7.86	7.84
8	10.67	9.84	9.75
16	15.09	14.02	13.94
32	21.58	20.41	20.14
64	33.05	31.67	31.53
128	51.30	49.54	48.68
256	92.62	89.04	86.41

…rt shapes Cherry-picked PR #3470's last-2-days commits (aux instance codegen, MAX_M removal, gemm2 BM=64 A_scale fix, AITER_MXFP4_INTERMEDIATE gate, tests). Two adaptations for this branch: - The PR codegen references the aux kernels under aux/; our branch renamed that to moe_aux/ (commit fee65e2, NFS compatibility). Repoint the codegen kernel includes, the dispatch-header include path, and the header location to moe_aux/. - Add AUX_EXTRA_SHAPES (MiniMax NE256/H3072, dsv3 NE32, Kimi-K2 NE384, Qwen3.5 NE512/TOPK10/H4096) to the AUX (sort) codegen only -- not the HIP gemm codegen -- so the BM16 inline_quant sort instances exist for the FlyDSL-port shapes without emitting HIP gemm instances that needn't compile. bench_up_moe_v2 --shapes all: all 9 shapes 1.08-1.68x vs fly, cos 0.97-0.99, correct across M (8/32/128). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: fsx950223 <fsx950223@outlook.com>

coderfeli · 2026-06-12T07:32:25Z

@ColorsWind conflicts

… large M) The synthesized (forced / no-CSV) mxfp4 port pipeline pinned BM16 inline_quant at all M -- decode-tuned but a big loss at prefill (M>=2048 ~main+40%, since BM16 atomic is latency-bound). Pick BM by M instead: BM16 inline_quant+atomic for M<=1024, BM128 nonatomic for M>1024 (throughput-bound prefill). Only the Kimi/DSR SHAPES have the non-inline BM128 path codegen'd (threestage sort + quant + sort_scales), so the tier switch is gated to them; other shapes stay BM16. KIMI (test_moe_2stage, PR vs main): M=4096 750us vs 793 (PR +6%), M=8192 1361 vs 1374 (PR +1%) -- recovering the large-M win that forced-BM16 lost (was main +40% at 2048, ~tie-but-slow at 8192). Matches PR #3470 compare.md's decode-win / mid-valley / large-M-win profile. Thresholds are heuristic; production tunes via CSV. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: fsx950223 <fsx950223@outlook.com>

fsx950223 · 2026-06-17T05:48:59Z

Could the kernel support general weight and scale shuffle? And the performance of kernels which INTER=256 are not as good as INTER=512.

ftyghome · 2026-06-17T06:07:55Z

Hi @fsx950223

Thanks for the review!

Could the kernel support general weight and scale shuffle?

AITER's shuffle_weights already provides both guinterleave=True and guinterleave=False variants, controlled by the ATOM_MOE_GU_ITLV=1 environment variable. The DeepSeek-V4 serving recipe (recipes/DeepSeek-V4.md) enables this environment variable as well.

Our kernel performs the silu(gate) * up fusion within a wave, and the guinterleave layout is the most efficient format for this access pattern. Our understanding is that this is a generally supported shuffle format. If we are misunderstanding the intended use case here, please let us know.

And the performance of kernels which INTER=256 are not as good as INTER=512.

Would you mind sharing your performance results for INTER=512?

Our current kernel has been extensively tuned and validated on Kimi K2.5, but we would also be happy to investigate and improve performance for other shapes if there are gaps.

fsx950223 · 2026-06-17T06:17:28Z

Hi @fsx950223

Thanks for the review!

Could the kernel support general weight and scale shuffle?

AITER's shuffle_weights already provides both guinterleave=True and guinterleave=False variants, controlled by the ATOM_MOE_GU_ITLV=1 environment variable. The DeepSeek-V4 serving recipe (recipes/DeepSeek-V4.md) enables this environment variable as well.

Our kernel performs the silu(gate) * up fusion within a wave, and the guinterleave layout is the most efficient format for this access pattern. Our understanding is that this is a generally supported shuffle format. If we are misunderstanding the intended use case here, please let us know.

And the performance of kernels which INTER=256 are not as good as INTER=512.

Would you mind sharing your performance results for INTER=512?

Our current kernel has been extensively tuned and validated on Kimi K2.5, but we would also be happy to investigate and improve performance for other shapes if there are gaps.

This is current branch vs main branch moe performance compare

It shows qwen35_397b is slower than main moe with large M, which INTER is 256

…figs Tune the flydsl a4w4 port on the merged branch and deploy the results. tune_mxfp4_moe.py: - Tag emitted rows `_tag=mxfp4_moe` (was mxfp4_guinterleave) so they route to the flydsl port on this branch (mxfp4_guinterleave now selects randomflow's #3470). - Gate variants against a torch bf16 MoE reference (new torch_moe_ref) for ALL shapes instead of run_fly -- the legacy CK 2stages reference is a slow per-shape build and segfaults on dsv4. Correct a4w4 variants land ~0.95 cos vs this ref; GATE lowered 0.90 -> 0.85 for margin. aiter/configs/tuned_fmoe.csv: - Add 56 tuned mxfp4_moe rows (14 shapes x M=8/128/2048/8192). The auto-merge had dropped the port's tuned rows (0 mxfp4_moe before); this restores them. No clash with the #3470 mxfp4_guinterleave rows (disjoint shape keys). Picks: BM16 atomic (decode) -> BM32/BM64 cshuffle (mid) -> BM128 nonatomic (prefill). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: fsx950223 <fsx950223@outlook.com>

…e + dedup The 107 mxfp4_guinterleave rows the merge added to tuned_fmoe.csv (base had 0) are the dev-branch flydsl PORT's tuned rows -- the dev port used the "mxfp4_guinterleave" shuffle_kind. On this branch the flydsl port routes via "mxfp4_moe" (mxfp4_guinterleave now selects randomflow's HIP #3470), so retag them mxfp4_moe. Then dedup the mxfp4_moe rows by shape key keeping best-us, which folds the freshly-tuned 56 rows from the prior commit into the 107 (8 token tiers). Result: 107 mxfp4_moe rows, 0 duplicates, 0 guinterleave in tuned_fmoe.csv. randomflow's base rows (kimik2_5 model_config) are untouched. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: fsx950223 <fsx950223@outlook.com>

The flydsl a4w4 MoE port was named inconsistently: the dispatch and the tuned CSV `_tag` use "mxfp4_moe", but the dev branch's old port alias "mxfp4_guinterleave" survived in _MXFP4_PORT_KIND, the header/_port_kind comments, and 4 test sites. Because _port_kind = w1.shuffle_kind == _MXFP4_PORT_KIND # "mxfp4_guinterleave" a model setting only shuffle_kind="mxfp4_moe" (the natural, CSV-consistent way) got _port_kind=False -> the flydsl gemm backend was not auto-selected -> silent fallback to HIP gemm. It was masked by callers also setting gemm{1,2}_backend="flydsl" explicitly. Unify on "mxfp4_moe" as the single port kind: - _MXFP4_PORT_KIND="mxfp4_moe", _MXFP4_KINDS=("mxfp4_moe",) - fix the (inverted) header comment + the _port_kind comment - tests: shuffle_kind="mxfp4_moe"; drop the now-redundant explicit gemm{1,2}_backend from test_moe_2stage so it verifies shuffle_kind-alone auto-select AND hits the tuned mxfp4_moe CSV rows (BM64/cshuffle) instead of the M-adaptive synthesis fallback. "mxfp4_guinterleave" now refers exclusively to randomflow's HIP a4w4 (#3470). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: fsx950223 <fsx950223@outlook.com>

* [PERF] MXFP4 (a4w4) MoE backend for gfx950 (#3470) --------- Co-authored-by: ColorsWind <14761584+ColorsWind@users.noreply.github.com> Co-authored-by: Zesen Liu <ftyghome@gmail.com> Co-authored-by: GnSight <ftyg@live.com> Co-authored-by: Wulley <bernard.liu@amd.com> Co-authored-by: who who who <fsx950223@outlook.com> Co-authored-by: zhutaoyu <zhutaoyu97@gmail.com> Co-authored-by: Claude Opus 4 (1M context) <noreply@anthropic.com> Co-authored-by: zhangnju <zhangnju@126.com> Co-authored-by: Bernard-Liu <bernaliu@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>

ColorsWind added 4 commits June 1, 2026 15:26

feat(mxfp4_moe): MoE sort/quant/scatter-reduce aux kernels

980f940

Add module_moe_mxfp4_aux for the MXFP4 (a4w4) MoE backend: fused token sort+quant, 3-stage sort, scale sort/shuffle, and weighted bf16 scatter-reduce (gfx950, Kimi-K2.5 / DSR shapes).

ColorsWind requested a review from a team June 1, 2026 15:52

srinivamd mentioned this pull request Jun 3, 2026

fix(moe): route SwiGLU MXFP4 unshuffled weights to CK-Tile instead of CK2stages #3518

Merged

ftyghome mentioned this pull request Jun 3, 2026

perf: add preshuffle logic for HIP fused_moe ROCm/ATOM#1057

Closed

coderfeli reviewed Jun 5, 2026

View reviewed changes

coderfeli changed the base branch from main to dev/randomflow_pr June 7, 2026 01:38

ColorsWind added 6 commits June 10, 2026 12:57

ColorsWind and others added 2 commits June 10, 2026 18:33

test(mxfp4_moe): cleanup test code

ac35a54

Merge branch 'dev/randomflow_pr' into rf-mxfp4-moe

57e7450

ftyghome added 2 commits June 12, 2026 12:52

refactor(mxfp4_moe): respect guinterleave shuffle

39e2ed7

refactor: gate by GateMode.INTERLEAVE, not by custom tag.

d6ede45

coderfeli merged commit fa0763b into ROCm:dev/randomflow_pr Jun 15, 2026
2 checks passed

fsx950223 mentioned this pull request Jun 18, 2026

Port flydsl a4w4 MoE onto randomflow_pr (coexist with HIP #3470) #3788

Merged

coderfeli pushed a commit that referenced this pull request Jun 18, 2026

Port flydsl a4w4 MoE onto randomflow_pr (coexist with HIP #3470) (#3788)

46ecc44

jhinpan mentioned this pull request Jun 19, 2026

[Issue]: MXFP4 MoE low MFU at large shapes and long latency at small tokens ROCm/FlyDSL#708

Closed

ColorsWind mentioned this pull request Jun 21, 2026

[RadeonFlow] flydsl mxfp4 a4w4 MoE: bring gemm1/gemm2 to parity with the HIP backend #3828

Merged

fsx950223 mentioned this pull request Jun 23, 2026

[draft] flydsl mxfp4 gemm2: pipeline fp4-out epilog LDS reads (lgkmcnt -68%, … #3840

Merged

ColorsWind mentioned this pull request Jul 1, 2026

[RadeonFlow] MXFP4 (a4w4) MoE backend for gfx950 #3832

Merged

1 task

benenzhu mentioned this pull request Jul 11, 2026

[AMD][ATOM] KIMI fp4 MI355 update docker SemiAnalysisAI/InferenceX#2132

Merged

Uh oh!

Conversation

ColorsWind commented Jun 1, 2026

[PERF] mxfp4_moe — a native MXFP4 (a4w4) MoE backend for gfx950

Motivation

Technical Details

Several pipelines, dispatched per token-count M

Test Plan

Correctness

Performance (end-to-end, mxfp4_moe vs FlyDSL)

Uh oh!

github-actions Bot commented Jun 1, 2026

🏷️ CI Guide

Uh oh!

benenzhu commented Jun 2, 2026

Uh oh!

coderfeli commented Jun 3, 2026

Uh oh!

ftyghome commented Jun 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

lalala-sh commented Jun 3, 2026

Uh oh!

benenzhu commented Jun 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

coderfeli Jun 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ftyghome commented Jun 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ftyghome commented Jun 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Technical detail

Test Plan

Accuracy

Performance

Detailed Performance

Interactivity, tok/s/user

TPOT, ms

Throughput/GPU, tok/s

E2E latency, s

Uh oh!

coderfeli commented Jun 12, 2026

Uh oh!

Uh oh!

ftyghome commented Jun 3, 2026 •

edited

Loading

benenzhu commented Jun 3, 2026 •

edited

Loading

coderfeli Jun 5, 2026 •

edited

Loading

ftyghome commented Jun 11, 2026 •

edited

Loading

ftyghome commented Jun 12, 2026 •

edited

Loading

fsx950223 commented Jun 17, 2026 •

edited

Loading